"""
OpenCTI Connector for importing threat intelligence data into the Knowledge Graph.
"""
import os
import logging
import time
from datetime import datetime
from typing import Dict, List, Any, Optional
import json

from neo4j import GraphDatabase
from pycti import OpenCTIApiClient

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class OpenCTIConnector:
    """
    Connector for importing data from OpenCTI into the Knowledge Graph.
    """
    
    def __init__(self, neo4j_uri: str, neo4j_user: str, neo4j_password: str,
                 opencti_url: str, opencti_token: str):
        """
        Initialize the OpenCTI Connector.
        
        Args:
            neo4j_uri (str): Neo4j database URI
            neo4j_user (str): Neo4j username
            neo4j_password (str): Neo4j password
            opencti_url (str): OpenCTI API URL
            opencti_token (str): OpenCTI API token
        """
        # Initialize Neo4j connection
        try:
            self.neo4j_driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
            self.neo4j_driver.verify_connectivity()
            logger.info("Connected to Neo4j database")
        except Exception as e:
            logger.error(f"Failed to connect to Neo4j: {str(e)}")
            self.neo4j_driver = None
            raise
        
        # Initialize OpenCTI client
        try:
            self.opencti_client = OpenCTIApiClient(opencti_url, opencti_token)
            logger.info("Connected to OpenCTI API")
        except Exception as e:
            logger.error(f"Failed to connect to OpenCTI: {str(e)}")
            if self.neo4j_driver:
                self.neo4j_driver.close()
            self.opencti_client = None
            raise
    
    def close(self):
        """Close all connections."""
        if self.neo4j_driver:
            self.neo4j_driver.close()
            logger.info("Neo4j connection closed")
    
    def import_stix_objects(self, limit: int = 100, types: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        Import STIX objects from OpenCTI into the Knowledge Graph.
        
        Args:
            limit (int): Maximum number of objects to import
            types (List[str], optional): List of STIX types to import
            
        Returns:
            Dict: Import statistics
        """
        if not self.opencti_client or not self.neo4j_driver:
            logger.error("Cannot import: connections not initialized")
            return {"error": "Connections not initialized", "imported": 0}
        
        # Default types if none provided
        if not types:
            types = ["threat-actor", "malware", "attack-pattern", "indicator", "vulnerability", "tool"]
        
        stats = {
            "total_imported": 0,
            "by_type": {},
            "errors": 0,
            "started_at": datetime.now().isoformat()
        }
        
        try:
            for stix_type in types:
                imported_count = self._import_stix_type(stix_type, limit)
                stats["by_type"][stix_type] = imported_count
                stats["total_imported"] += imported_count
                
            # Import relationships after entities
            relationship_count = self._import_relationships(limit)
            stats["by_type"]["relationship"] = relationship_count
            stats["total_imported"] += relationship_count
            
            stats["completed_at"] = datetime.now().isoformat()
            logger.info(f"Import completed: {stats['total_imported']} objects imported")
            return stats
            
        except Exception as e:
            logger.error(f"Error during import: {str(e)}")
            stats["error"] = str(e)
            stats["completed_at"] = datetime.now().isoformat()
            return stats
    
    def _import_stix_type(self, stix_type: str, limit: int) -> int:
        """
        Import STIX objects of a specific type.
        
        Args:
            stix_type (str): STIX type to import
            limit (int): Maximum number of objects to import
            
        Returns:
            int: Number of objects imported
        """
        logger.info(f"Importing {stix_type} objects (max {limit})")
        imported_count = 0
        
        # Fetch objects from OpenCTI
        try:
            # Different handling based on object type
            if stix_type == "threat-actor":
                objects = self.opencti_client.threat_actor.list(first=limit)
            elif stix_type == "malware":
                objects = self.opencti_client.malware.list(first=limit)
            elif stix_type == "attack-pattern":
                objects = self.opencti_client.attack_pattern.list(first=limit)
            elif stix_type == "indicator":
                objects = self.opencti_client.indicator.list(first=limit)
            elif stix_type == "vulnerability":
                objects = self.opencti_client.vulnerability.list(first=limit)
            elif stix_type == "tool":
                objects = self.opencti_client.tool.list(first=limit)
            else:
                logger.warning(f"Unsupported STIX type: {stix_type}")
                return 0
            
            # Process fetched objects
            if objects:
                with self.neo4j_driver.session() as session:
                    for obj in objects["edges"]:
                        node = obj["node"]
                        # Create or update node in Neo4j
                        result = session.run("""
                            MERGE (e:Entity {id: $id})
                            ON CREATE SET 
                                e.name = $name,
                                e.type = $type,
                                e.description = $description,
                                e.created = $created,
                                e.modified = $modified,
                                e.stix_id = $stix_id,
                                e.source = 'OpenCTI',
                                e.created_at = datetime()
                            ON MATCH SET 
                                e.name = $name,
                                e.type = $type,
                                e.description = $description,
                                e.modified = $modified,
                                e.updated_at = datetime()
                            RETURN e
                        """, {
                            "id": node["id"],
                            "name": node.get("name", "Unnamed"),
                            "type": stix_type,
                            "description": node.get("description", ""),
                            "created": node.get("created_at", ""),
                            "modified": node.get("updated_at", ""),
                            "stix_id": node.get("standard_id", "")
                        })
                        
                        imported_count += 1
                        if imported_count % 10 == 0:
                            logger.info(f"Imported {imported_count} {stix_type} objects")
            
            return imported_count
            
        except Exception as e:
            logger.error(f"Error importing {stix_type}: {str(e)}")
            raise
    
    def _import_relationships(self, limit: int) -> int:
        """
        Import relationships between STIX objects.
        
        Args:
            limit (int): Maximum number of relationships to import
            
        Returns:
            int: Number of relationships imported
        """
        logger.info(f"Importing relationships (max {limit})")
        imported_count = 0
        
        try:
            # Fetch relationships from OpenCTI
            relationships = self.opencti_client.stix_relation.list(first=limit)
            
            if relationships:
                with self.neo4j_driver.session() as session:
                    for rel in relationships["edges"]:
                        relationship = rel["node"]
                        
                        # Skip if source or target is missing
                        if not relationship.get("from") or not relationship.get("to"):
                            continue
                        
                        # Create relationship in Neo4j
                        result = session.run("""
                            MATCH (source:Entity {stix_id: $source_id})
                            MATCH (target:Entity {stix_id: $target_id})
                            MERGE (source)-[r:`RELATES_TO` {id: $id}]->(target)
                            ON CREATE SET 
                                r.type = $type,
                                r.description = $description,
                                r.created = $created,
                                r.modified = $modified,
                                r.stix_id = $stix_id,
                                r.source = 'OpenCTI',
                                r.detection_method = 'opencti_import',
                                r.created_at = datetime()
                            ON MATCH SET 
                                r.type = $type,
                                r.description = $description,
                                r.modified = $modified,
                                r.updated_at = datetime()
                            RETURN r
                        """, {
                            "id": relationship["id"],
                            "source_id": relationship["from"]["standard_id"],
                            "target_id": relationship["to"]["standard_id"],
                            "type": relationship.get("relationship_type", "relates-to"),
                            "description": relationship.get("description", ""),
                            "created": relationship.get("created_at", ""),
                            "modified": relationship.get("updated_at", ""),
                            "stix_id": relationship.get("standard_id", "")
                        })
                        
                        imported_count += 1
                        if imported_count % 10 == 0:
                            logger.info(f"Imported {imported_count} relationships")
            
            return imported_count
            
        except Exception as e:
            logger.error(f"Error importing relationships: {str(e)}")
            raise

    def run_incremental_import(self, interval_hours: int = 24, types: Optional[List[str]] = None):
        """
        Run an incremental import job at regular intervals.
        
        Args:
            interval_hours (int): Number of hours between imports
            types (List[str], optional): List of STIX types to import
        """
        while True:
            try:
                logger.info(f"Starting incremental import")
                stats = self.import_stix_objects(limit=500, types=types)
                logger.info(f"Incremental import completed: {stats['total_imported']} objects imported")
                
                # Sleep until next import
                next_import = datetime.now().timestamp() + (interval_hours * 3600)
                logger.info(f"Next import scheduled in {interval_hours} hours")
                time.sleep(interval_hours * 3600)
                
            except Exception as e:
                logger.error(f"Error in incremental import: {str(e)}")
                # Wait a bit before retrying
                time.sleep(300) 